import os
import sys
import nbformat
from nbconvert import PythonExporter, MarkdownExporter

def convert_notebooks():
    # 创建输出文件夹
    if not os.path.exists('py'):
        os.makedirs('py')
    if not os.path.exists('md'):
        os.makedirs('md')

    way = 'py' if len(sys.argv)<2 or sys.argv[1]=='py' else 'md'

    # 获取当前目录下的所有.ipynb文件
    for filename in os.listdir('.'):
        if filename.endswith('.ipynb'):
            # 读取.ipynb文件
            with open(filename, 'r', encoding='utf-8') as f:
                notebook_content = nbformat.read(f, as_version=4)

            # 转换为.py文件
            if way == "py":
                python_exporter = PythonExporter()
                (python_code, _) = python_exporter.from_notebook_node(notebook_content)
                py_filename = os.path.join('py', filename.replace('.ipynb', '.py'))
                with open(py_filename, 'w', encoding='utf-8') as f:
                    f.write(python_code)
                    print(f"Converted {filename} to {py_filename}")
            elif way == "md":
            # 转换为.md文件
                markdown_exporter = MarkdownExporter()
                (markdown_code, _) = markdown_exporter.from_notebook_node(notebook_content)
                md_filename = os.path.join('md', filename.replace('.ipynb', '.md'))
                with open(md_filename, 'w', encoding='utf-8') as f:
                    f.write(markdown_code)
                    print(f"Converted {filename} to {md_filename}")

                

if __name__ == '__main__':
    convert_notebooks()